'''
引入大智慧股票数据csv和bww_macd1数据文件,bwwmacd数据导入为data_signals,

买入： 根据里面的m5,利用内置均线公式计算金叉
       附加判断低迷期低于股票总数的75%, 程序设定小于-500 .

止损：  指标多头状态，跌破买入价执行
        macd 小于0, 跌破买入价 1% 否则 3%。

仓位控制：  macd 小于0, 买入1份, 大于0-买入2份
'''
import backtrader as bt  
import pandas as pd
from datetime import datetime
import matplotlib.pyplot as plt

from backtrader_plotting import Bokeh
from backtrader_plotting.schemes import Tradimo

# 读取csv
data_stock = pd.read_csv('法宝回测/data/sh1.csv') # 股票日数据
data_signals = pd.read_csv('法宝回测/data/bwwmacd.csv') # 股票指标数据(大智慧导出)

# 提取需要列
selected_columns = ['日期', '开盘', '最高', '最低', '收盘', '成交量' ]
data_stock = data_stock[selected_columns]
selected_columns = ['日期','MA1','MA2','MA3','S' ]
data_signals = data_signals[selected_columns]

# 改变列名称
data_stock.columns = ['date', 'open', 'high', 'low', 'close', 'volume']
data_signals.columns = ['date', 'M5', 'M10', 'M20','MA']

# 日期列转换时间类型
data_stock['date'] = pd.to_datetime(data_stock['date'])
data_signals['date'] = pd.to_datetime(data_signals['date'])

# 时间索引
data_stock.index = pd.to_datetime(data_stock['date'])
data_signals.index = pd.to_datetime(data_signals['date'])

# backtrader投喂数据格式
data_stock = data_stock.iloc[:, :6]
data_stock['openinterest'] = 0

#==============================================================================
# 自定义数据源，这里仅展示框架，具体实现需要根据你的CSV文件格式进行调整  
class MyStockData(bt.feeds.PandasData):  # MyStockData 类定义了如何从Pandas DataFrame中提取股票数据。
    # params 元组指定了日期、开盘价、最高价、最低价、收盘价和交易量（如果有的话）在DataFrame中的列索引。
    params = (  
        ('datetime', 0),  # 日期在DataFrame中的列索引  
        ('open', 1),      # 开盘价在DataFrame中的列索引  
        ('high', 2),      # 最高价在DataFrame中的列索引  
        ('low', 3),       # 最低价在DataFrame中的列索引  
        ('close', 4),     # 收盘价在DataFrame中的列索引  
        ('volume', 5),    # 交易量在DataFrame中的列索引（如果有的话）  
        ('openinterest', -1),    # 交易量在DataFrame中的列索引（如果有的话）  
    )  
class MySignalData(bt.feeds.PandasData):  # MySignalData 类用于处理信号数据。
    lines = ('M5','M10','M20','MA',)  # lines = ('signal',) 是一个元组的示例，
    # 指定了日期和信号在DataFrame中的列索引。
    params = (  
        ('datetime', 0),  # 日期在DataFrame中的列索引
        ('M5', 1), # 假设'M5'是某种信号，这里重命名以避免混淆  
        ('M10', 2), # 同上  
        ('M20', 3), # 同上 
        ('MA', 4), # 同上 
    )  
    
#===========策略策略策略策略策略策略策略策略策略策略策略======================================

# 在策略中， self.datas[0] 访问股票数据，通过 self.datas[1] 访问信号数据。
class MyStrategy(bt.Strategy):  
    # params这个参数用于控制是否打印日志到控制台。
    params = (  
        ('printlog', True),  
        #('stop_loss', 0.03),  # 止损比例，比如3%
        ("ma", 5), 
        ("fast_period", 1),  # 短期（5日）移动平均线周期
        ("slow_period", 5),  # 长期（30日）移动平均线周期
        
        ("macd_fast", 12),  # MACD快线周期
        ("macd_slow", 26),  # MACD慢线周期
        ("macd_signal", 9),  # MACD信号线周期
    )  # 只包含一个 两元素的子元组('printlog', True) 的元组。可以通过索引来访问这个子元组及其内部的元素。

    # 用于记录日志的函数
    def log(self, txt, dt=None):  
        if self.params.printlog:  
            dt = dt or self.datas[0].datetime.datetime(0)  
            print(f'{dt}, {txt}')  

    def __init__(self):  
        # 初始化策略  
        self.dataclose = self.datas[0].close  
        
        # 记录买入价格
        self.buy_price = None
        
        sma5 = bt.ind.SMA(period=self.p.ma)  # fast moving average
        # 初始化短期和长期移动均线指标
        self.fast_sma = bt.indicators.SimpleMovingAverage(self.datas[1].M5, period=self.params.fast_period)
        self.slow_sma = bt.indicators.SimpleMovingAverage(self.datas[1].M5, period=self.params.slow_period)
        self.crossover = bt.ind.CrossOver(self.fast_sma, self.slow_sma)  # crossover signal
        
        # 初始化MACD指标
        self.macd = bt.indicators.MACDHisto(self.datas[0],
                                           period_me1=self.params.macd_fast,
                                           period_me2=self.params.macd_slow,
                                           period_signal=self.params.macd_signal)
        # 检查MACD金叉和死叉
        self.macd_line = self.macd.lines.macd  # MACD线（快线-慢线）
        self.signal_line = self.macd.lines.signal  # 信号线
        self.buysize = 0
        self.stop_loss=0

    #===========================================NEXT===========================================
    def next(self):  
        # 计算止损价格
        if self.macd_line < 0:
            self.stop_loss=0.01
        else:
            self.stop_loss=0.03
        stop_price = self.buy_price * (1 - self.stop_loss) if self.buy_price else 0
        
        # 简单示例：根据信号进行买入或卖出操作, 如果datasignal的当前值（self.datasignal[0]）...
        if not self.position:  # 如果没有持仓 
            if self.crossover > 0 and self.fast_sma <-500: # or self.fast_sma > 0 ):# and self.data1M5 < self.data1MA:  # 如果信号为1，买入  
                print('买入')
                if self.macd_line < 0:
                    
                    self.buysize=1
                    self.buy(size=self.buysize)
                elif self.macd_line > 0:
                    self.buysize=2
                    self.buy(size=self.buysize)
                    
                self.log('BUY CREATE, %.2f' % self.dataclose[0])  
                # self.buy(size=2)  
                # 更新买入价格
                self.buy_price = self.data.close[0]
                
        else:
            
            if  self.crossover < 0:    # 如果信号为-1，卖出
                print('卖出')  
                self.log('SELL CREATE, %.2f' % self.dataclose[0])  
                self.sell(size=self.buysize) 
            elif self.fast_sma > self.slow_sma and self.data.close[0] < stop_price:
                print('止损') 
                self.log('SELL CREATE, %.2f' % self.dataclose[0])  
                self.sell(size=self.buysize) 
                
#===============================================================================
# 创建Cerebro引擎  
cerebro = bt.Cerebro()  

# # 设置回测时间范围
start_date = datetime(2023, 5, 1)
end_date = datetime(2024, 6, 14)

# 添加数据源  
datafeed_stock = MyStockData(dataname=data_stock, fromdate=start_date, todate=end_date)  
datafeed_signal = MySignalData(dataname=data_signals, fromdate=start_date, todate=end_date)  

# 将数据源添加到cerebro  
cerebro.adddata(datafeed_stock, name="000001")  
cerebro.adddata(datafeed_signal, name="bwwmacd")  

# 添加策略  
cerebro.addstrategy(MyStrategy)  

# 设置初始资金和手续费
start_cash = 8000
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.001)

# 添加策略分析指标
cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='tradeanalyzer')  
cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='annualReturn')  
cerebro.addanalyzer(bt.analyzers.Returns, _name='annualizedReturns', tann=252)  # 使用 'annualizedReturns' 代替 '_Returns'  
cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')  
cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpeRatio')  # 使用 'sharpeRatio' 代替 'sharpe'  
cerebro.addanalyzer(bt.analyzers.Returns, _name='totalReturns')  # 使用 'totalReturns' 代替重复的 'returns'  
# 假设 'TimeReturn' 是有效的分析器  
cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='timeReturn')  # 假设这是有效的

# 运行回测  
results=cerebro.run()  

# 输出结果  
#获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")#计算胜率

# 打印分析器输出结果
total_trades = results[0].analyzers.tradeanalyzer.get_analysis()['total']['total']
won_trades = results[0].analyzers.tradeanalyzer.get_analysis()['won']['total']
win_rate = (won_trades / total_trades) * 100 if total_trades > 0 else 0
print('总交易次数:', total_trades)
print('盈利次数:', won_trades)
print('胜率%:', win_rate)

# print(f"初始资金: {start_cash}\n回测期间:{start_date.strftime('%Y-%m-%d')} : {end_date.strftime('%Y-%m-%d')}")
#print('年度汇报:', results[0].analyzers.annualReturn.get_analysis())
print('年化收益%:', results[0].analyzers.annualizedReturns.get_analysis()['rnorm100'])
print('最大回撤比例%:', results[0].analyzers.drawdown.get_analysis().max.drawdown)
print('夏普比率:', results[0].analyzers.sharpeRatio.get_analysis()['sharperatio'])
print('累计收益：', results[0].analyzers.totalReturns.get_analysis()['rtot'])

print(f'最后投资金额：{round(cerebro.broker.getvalue(), 2)}')
# ==================================================================================
plotconfig = {
    'id:ind#0': dict(
        subplot=True,
    ),
}
b = Bokeh(style='line', scheme=Tradimo(),plotconfig=plotconfig)
cerebro.plot(b)